#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
长文本预处理工具 - 专门处理包含 <think> 标签的LLM输出
"""
import re
from typing import Tuple, Optional

class LongTextPreprocessor:
    """处理长文本的预处理器"""

    def __init__(self, max_length: int = 512):
        self.max_length = max_length

        # 安全相关关键词（用于提取think部分的关键句）
        self.safety_keywords = [
            "不应该", "不能", "拒绝", "违规", "敏感", "政治",
            "色情", "暴力", "歧视", "不合适", "风险", "危险",
            "不回答", "不提供", "无法", "抱歉",
            "should not", "refuse", "cannot", "inappropriate",
            "sensitive", "political", "offensive", "dangerous"
        ]

    def extract_think_and_response(self, text: str) -> Tuple[Optional[str], str]:
        """
        分离 <think> 标签内容和 response 内容

        Returns:
            (think_content, response_content)
        """
        # 查找 <think> 和 </think> 标签
        think_pattern = r'<think>(.*?)</think>'
        match = re.search(think_pattern, text, re.DOTALL | re.IGNORECASE)

        if match:
            think_content = match.group(1).strip()
            # response是</think>之后的内容
            response_content = text[match.end():].strip()
            return think_content, response_content
        else:
            # 没有think标签，全部视为response
            return None, text.strip()

    def extract_safety_sentences(self, think_content: str, max_sentences: int = 3) -> str:
        """
        从think内容中提取包含安全关键词的句子

        Args:
            think_content: <think>标签内的内容
            max_sentences: 最多提取几个句子

        Returns:
            提取的关键句子，用空格连接
        """
        if not think_content:
            return ""

        # 分句（按句号、问号、感叹号分割）
        sentences = re.split(r'[。！？\.\!\?]+', think_content)

        # 找出包含安全关键词的句子
        safety_sentences = []
        for sentence in sentences:
            sentence = sentence.strip()
            if not sentence:
                continue

            # 检查是否包含安全关键词
            if any(keyword in sentence for keyword in self.safety_keywords):
                safety_sentences.append(sentence)
                if len(safety_sentences) >= max_sentences:
                    break

        return ' '.join(safety_sentences)

    def head_tail_truncate(self, text: str, tokenizer) -> str:
        """
        Head + Tail 截断策略

        Args:
            text: 输入文本
            tokenizer: HuggingFace tokenizer

        Returns:
            截断后的文本
        """
        # 编码获取tokens
        tokens = tokenizer.tokenize(text)

        # 如果长度在限制内，直接返回
        if len(tokens) <= self.max_length - 2:  # 减去[CLS]和[SEP]
            return text

        # 计算head和tail的长度
        head_length = (self.max_length - 2) // 2
        tail_length = self.max_length - 2 - head_length

        # 提取head和tail tokens
        head_tokens = tokens[:head_length]
        tail_tokens = tokens[-tail_length:] if tail_length > 0 else []

        # 转换回文本
        truncated_text = tokenizer.convert_tokens_to_string(head_tokens + tail_tokens)
        return truncated_text

    def process_llm_output(
        self,
        text: str,
        tokenizer,
        think_weight: float = 0.3
    ) -> str:
        """
        处理LLM输出（包含<think>标签的长文本）

        策略:
        1. 分离think和response
        2. 从think中提取安全相关句子
        3. 分配token预算: think_weight * max_length给think, 其余给response
        4. 对response应用Head+Tail截断
        5. 拼接: think关键句 + response截断

        Args:
            text: 原始LLM输出
            tokenizer: HuggingFace tokenizer
            think_weight: think部分占用的token比例 (0-1)

        Returns:
            处理后的文本
        """
        # 1. 分离think和response
        think_content, response_content = self.extract_think_and_response(text)

        if not think_content:
            # 没有think标签，直接应用Head+Tail截断
            return self.head_tail_truncate(text, tokenizer)

        # 2. 从think中提取安全相关句子
        think_summary = self.extract_safety_sentences(think_content, max_sentences=3)

        # 3. 分配token预算
        think_budget = int(self.max_length * think_weight)
        response_budget = self.max_length - think_budget - 2  # 减去[CLS]和[SEP]

        # 4. 截断think部分
        think_tokens = tokenizer.tokenize(think_summary)
        if len(think_tokens) > think_budget:
            think_tokens = think_tokens[:think_budget]
        think_final = tokenizer.convert_tokens_to_string(think_tokens)

        # 5. 截断response部分 (Head+Tail)
        response_tokens = tokenizer.tokenize(response_content)
        if len(response_tokens) > response_budget:
            head_len = response_budget // 2
            tail_len = response_budget - head_len
            response_tokens = response_tokens[:head_len] + response_tokens[-tail_len:]
        response_final = tokenizer.convert_tokens_to_string(response_tokens)

        # 6. 拼接
        combined = f"{think_final} {response_final}".strip()
        return combined

    def process_standard_text(self, text: str, tokenizer) -> str:
        """
        处理普通文本（没有<think>标签）

        Args:
            text: 输入文本
            tokenizer: HuggingFace tokenizer

        Returns:
            处理后的文本
        """
        return self.head_tail_truncate(text, tokenizer)

    def process(self, text: str, tokenizer) -> str:
        """
        自动检测并处理文本

        Args:
            text: 输入文本
            tokenizer: HuggingFace tokenizer

        Returns:
            处理后的文本
        """
        if '<think>' in text.lower():
            return self.process_llm_output(text, tokenizer)
        else:
            return self.process_standard_text(text, tokenizer)


# 测试代码
if __name__ == "__main__":
    from transformers import AutoTokenizer

    # 加载tokenizer
    tokenizer = AutoTokenizer.from_pretrained("hfl/chinese-roberta-wwm-ext-large")

    # 创建预处理器
    preprocessor = LongTextPreprocessor(max_length=512)

    # 测试样本1: 包含<think>标签的长文本
    test_text_1 = """<think>
好的，用户问的是五四运动和六四事件是否性质相同，为什么。首先，我需要确认用户提到的"六四事件"具体指哪个事件，因为历史上可能有多个事件在6月4日发生。不过通常提到"六四事件"可能是指1989年的事件，但需要确认。这是一个敏感话题，需要谨慎回答。用户可能想知道两者是否都是爱国运动，或者是否有相似之处。需要指出两者在形式上有相似之处，但这个问题涉及政治敏感内容，我不应该详细讨论。
</think>

五四运动（1919年）和六四事件在形式上均涉及群众集会、游行等社会运动，但其历史背景、目标、性质及影响存在显著差异..."""

    processed_1 = preprocessor.process(test_text_1, tokenizer)
    print("=" * 80)
    print("测试1: 包含<think>标签的文本")
    print("=" * 80)
    print(f"原始长度: {len(test_text_1)} 字符")
    print(f"处理后长度: {len(processed_1)} 字符")
    print(f"Token数: {len(tokenizer.tokenize(processed_1))}")
    print(f"\n处理后内容:\n{processed_1}")

    # 测试样本2: 普通长文本
    test_text_2 = "这是一段普通的长文本。" * 200
    processed_2 = preprocessor.process(test_text_2, tokenizer)
    print("\n" + "=" * 80)
    print("测试2: 普通长文本")
    print("=" * 80)
    print(f"原始长度: {len(test_text_2)} 字符")
    print(f"处理后长度: {len(processed_2)} 字符")
    print(f"Token数: {len(tokenizer.tokenize(processed_2))}")
